import numpy as np
from mindquantum.circuit import decompose_single_term_time_evolution
from openfermion.ops import QubitOperator
from mindquantum import Circuit, UN
import mindquantum.gate as G
from mindquantum import Hamiltonian
from mindquantum.nn import generate_pqc_operator
from mindspore import Tensor
import matplotlib.pyplot as plt


def generate_random_sk(n, seed=42):
    np.random.seed(seed)
    node = {}
    for i in range(n):
        for j in range(i + 1, n):
            node[(i, j)] = np.random.choice([-1, 1])
    return node


def generate_hamiltonian(model, gate='ZZ'):
    def generate_gate(gate, poi1, poi2):
        return '{}{} {}{}'.format(gate[0], poi1, gate[1], poi2)

    q = QubitOperator()
    for (n1, n2), w in model.items():
        q += QubitOperator(generate_gate(gate, n1, n2), w * 1.0)
    return q


def generate_c_ansatz(model, para_name='', gate='ZZ'):
    def generate_gate(gate, poi1, poi2):
        return '{}{} {}{}'.format(gate[0], poi1, gate[1], poi2)

    c = Circuit()
    for (n1, n2), w in model.items():
        sub = decompose_single_term_time_evolution(
            QubitOperator(generate_gate(gate, n1, n2)), {para_name: w})
        c += sub
    return c


def generate_b_ansatz(n_qubits, gate='X', para_name=''):
    gates = {'X': G.RX, 'Y': G.RY, 'Z': G.RZ}
    c = Circuit()
    for i in range(n_qubits):
        c += gates[gate]({para_name: 1}).on(i)
    return c


def show_amp(state):
    amp = np.diag(state.dot(np.conj(state.T)))
    n_qubits = int(np.log2(len(amp)))
    labels = [bin(i)[2:].zfill(n_qubits) for i in range(len(amp))]
    plt.bar(labels, amp)
    plt.xticks(rotation=45)
    plt.show()


def update_bar(num, bars, bardata):
    for i, b in enumerate(bars):
        b.set_height(bardata[num, i])


def fun(para, pqc):
    a = Tensor(np.array([[]]).astype(np.float32))
    b = Tensor(np.array(para).astype(np.float32))
    f, _, g = pqc(a, b)
    return float(f.asnumpy()[0, 0]), g.asnumpy()[0, 0]


def fit_data(data):
    data_x = []
    data_y = []
    for x, f in data:
        data_x.append(x)
        data_y.append(f)
    data_x = np.array(data_x)
    data_y = np.array(data_y)
    data_x = data_x.T
    grad = []
    for x in data_x:
        g, _ = np.polyfit(x, data_y, 1)
        grad.append(g)
    return np.array(grad)


def find_min(pqc, ansatz_para):
    n = 500
    k = 20
    x0 = np.array([0] * len(ansatz_para)).astype(np.float32)
    for i in range(0, len(x0), 2):
        x0[i] = np.pi / 4
    r = 0.1
    alpha = 1.1
    delta = 0.2
    cta = 1.1
    A = 1
    eps = 0.001
    L = []
    x = x0 * 1
    m = 0
    step = 0
    while step + k < n:
        f = fun(x, pqc)[0]
        L.append((x, f))
        delta_p = delta / (m + 1)**cta
        for i in range(k):
            x_p = np.random.uniform(x - delta_p, x + delta_p)
        L.append((x_p, fun(x_p, pqc)[0]))
        L_p = []
        for x_p, f_p in L[-k:]:
            if ((x_p - x)**2).sum()**0.5 < delta_p:
                L_p.append((x_p, f_p))
        if len(L_p) > 1:
            g = fit_data(L_p)
            r_p = r / (m + 1 + A)**alpha
            if r_p * (g**2).sum()**0.5 < eps:
                return x
            x = x - r_p * g
            m = m + 1
        input((x, f))


if __name__ == '__main__':
    from scipy.optimize import minimize
    from mindquantum.nn import generate_evolution_operator
    from matplotlib import animation

    p = 2
    n_qubits = 5
    a = 1
    c = Circuit()
    model = {
        (0, 1): 1 * a,
        (1, 2): 1 * a,
        (2, 3): 1 * a,
        (0, 3): 1 * a,
        (0, 4): 1 * a
    }
    for i in range(p):
        c += generate_c_ansatz(model, 'g' + str(i), 'XX')
        c += generate_b_ansatz(n_qubits, 'Z', 'b' + str(i))

    ansatz_para = c.parameter_resolver().para_name
    ham = Hamiltonian(generate_hamiltonian(model, 'XX'))
    pqc = generate_pqc_operator([], ansatz_para, c, ham)
    # find_min(pqc, ansatz_para)
    params = []
    # res = minimize(
    #     fun,
    #     np.random.normal(size=len(ansatz_para))*0.01,
    #     # np.zeros(len(ansatz_para)),
    #     args=(pqc, ),
    #     method='bfgs',
    #     jac=True,
    #     options={'disp': True},
    #     callback=lambda x: params.append(x))
    params = np.array(params)
    params = np.load('params.npy')
    # final_para = res.x
    evol = generate_evolution_operator(ansatz_para, c + UN(G.H, n_qubits))
    amps = []
    for final_para in params:
        state = evol(Tensor(final_para.astype(np.float32)))
        state = np.round(state.asnumpy(), 3)
        amp = np.diag(state.dot(np.conj(state.T)))
        amps.append(amp)
    amps = np.array(amps)
    labels = [bin(i)[2:].zfill(n_qubits) for i in range(2**n_qubits)]
    fig = plt.figure()
    ax = fig.add_subplot()
    bars = plt.bar(labels, amps[0])
    plt.xticks(rotation=45)
    ax.set_ylim(0, 0.4)
    anim = animation.FuncAnimation(fig,
                                   update_bar,
                                   len(amps),
                                   fargs=(bars, amps),
                                   interval=50)
    plt.show()
    # show_amp(state)
    # s = state[:, 0] + 1j * state[:, 1]
    # print(np.vdot(s, s))
